import psutil
import time
import os


def get_process_io_counters(pid):
    try:
        process = psutil.Process(pid)
        io_counters = process.io_counters()
        return io_counters
    except (psutil.NoSuchProcess, psutil.AccessDenied):
        print(f"无法获取进程 {pid} 的 I/O 计数器信息")
        return None


def get_process_io_speed(pid, interval=1):
    start_io_counters = get_process_io_counters(pid)
    time.sleep(interval)
    end_io_counters = get_process_io_counters(pid)

    if start_io_counters and end_io_counters:
        read_speed = (
            end_io_counters.read_bytes - start_io_counters.read_bytes
        ) / interval
        write_speed = (
            end_io_counters.write_bytes - start_io_counters.write_bytes
        ) / interval
        return read_speed, write_speed
    else:
        return None


# pid = 12345  # 将此值替换为要监控的进程 ID
pid=int(input("手动运行 io_speed 测试程序: ./io_read_write_test.sh \n输入pid:"))
read_speed, write_speed = get_process_io_speed(pid)

if read_speed is not None and write_speed is not None:
    print(f"进程 {pid} 的读速率: {read_speed} bytes/s")
    print(f"进程 {pid} 的写速率: {write_speed} bytes/s")
else:
    print(f"无法获取进程 {pid} 的读写速率")
